#include "PowerSpectrum.h"
#include "Int.h"

#include <fstream>
using std::ifstream;
#include <iostream>
#include <strstream>
using std::istrstream;
#include <string>
using std::string;
#include <map>
using std::map;
#include <vector>
using std::vector;

#include "TMV.h"
using tmv::Vector;
using tmv::Matrix;
#include "dbg.h"
using std::endl;

//#define NAME "mb"
// Add NAME to the output names (lnlike.NAME.dat and predmapsq.NAME.debug)

//#define GT4  
// Only use R > 4 for map^2 data

//#define DILLUT  1.10
// Assume dillution error, so multiply each gamma by DILLUT

//#define STARTAT 1530

//#define BMODE -1.0

std::ostream* dbgout=0;
bool XDEBUG=false;

const char* cmbfile = "/home/mjarvis/CMBQuintchainNew.dat";
//const char* cmbfile = "/home/mjarvis/FlatPrior.dat";

#ifdef GT4
const double map_minr = 4.0;
#else
const double map_minr = 1.0;
#endif
const double map_maxr = 40.0;
const double var_minr = 40.0;
const double var_maxr = 100.0;
const double singthresh = 0.01;
const int datafrac = 6; // take every nth line of the data

const double PI = 3.141592653589793;
const double DEG = PI/180.;
const double ARCSEC = DEG/3600.;
const double ARCMIN = DEG/60.;

#ifdef DILLUT
const double DILLUTSQ = DILLUT*DILLUT;
#endif

void ReadData(Vector<double>*& r, Vector<double>*& mean, Matrix<double>*& cov, 
    size_t& nmapout, const string& datecode)
{
  dbg<<"Start ReadData\n";

  const size_t nfields = 12;
  const string datafile = "/home/mjarvis/plots/m2.out.";
  const char* fieldcodes[nfields] = { 
    "A", "G", "H", "J", "K", "L", "M", "N", "Q", "R", "T", "X" };

  vector<vector<double> > rmap(nfields);
  vector<vector<double> > map(nfields);
  //vector<vector<double> > sigmap(nfields);
  vector<vector<double> > rvar(nfields);
  vector<vector<double> > var(nfields);
  //vector<vector<double> > sigvar(nfields);

  for(size_t k=0;k<nfields;k++) {
    string name = datafile + datecode + '.' + fieldcodes[k];
    ifstream fin(name.c_str());
    if (!fin) myerror("opening data file:  ",name.c_str());

    dbg<<"Reading file: "<<name<<endl;
    string line;
    int count=0;
    while (getline(fin,line)) {
      if (line[0] == '#') continue;
      count++;
      if (count < datafrac) continue;
      count = 0;
      istrstream linein(line.c_str());
      double r,mapsq,mxsq,junk,gamsq;
      double sigmap,siggamsq;
      linein >>r>>mapsq>>sigmap>>mxsq>>junk>>junk>>junk>>junk>>junk>>gamsq>>siggamsq;

#ifdef BMODE
      mapsq += (BMODE)*mxsq;
#endif

#ifdef DILLUT
      mapsq *= DILLUTSQ;
      gamsq *= DILLUTSQ;
      //sigmap *= DILLUTSQ;
      //siggamsq *= DILLUTSQ;
#endif

      if (r > map_minr && r < map_maxr) {
	rmap[k].push_back(r*ARCMIN);
	map[k].push_back(mapsq);
	//sigmap[k].push_back(sigmap);
      }
      if (r > var_minr && r < var_maxr) {
	rvar[k].push_back(r*ARCMIN);
	var[k].push_back(gamsq);
	//sigvar[k].push_back(sigvar);
      }
    }
    dbg<<"sizes = "<<rmap[k].size()<<", "<<rvar[k].size()<<endl;
    if (k>0) {
      assert(rmap[k].size() == rmap[0].size());
      assert(rvar[k].size() == rvar[0].size());
    }
  }
  size_t nmap = rmap[0].size();
  size_t nvar = rvar[0].size();
  size_t ntot = nmap + nvar;
  dbg<<"nmap = "<<nmap<<", nvar = "<<nvar<<endl;

  r = new Vector<double>(ntot,0.);
  mean = new Vector<double>(ntot,0.);
  cov = new Matrix<double>(ntot,ntot,0.);

  for(size_t k=0;k<nfields;k++) {
    for(size_t i=0;i<nmap;i++) {
      (*r)[i] += rmap[k][i];
      (*mean)[i] += map[k][i];
      for(size_t j=0;j<=i;j++) {
	(*cov)(i,j) += map[k][i] * map[k][j];
      }
    }
    for(size_t i=0;i<nvar;i++) {
      (*r)[nmap+i] += rvar[k][i];
      (*mean)[nmap+i] += var[k][i];
      for(size_t j=0;j<=i;j++) {
	(*cov)(nmap+i,nmap+j) += var[k][i] * var[k][j];
      }
      for(size_t j=0;j<nmap;j++) {
	(*cov)(nmap+i,j) += var[k][i]*map[k][j];
      }
    }
  }
  for(size_t i=1;i<ntot;i++) cov->diag(i) = cov->diag(-i);

  *r /= double(nfields);
  *mean /= double(nfields);
  *cov -= double(nfields)*OuterProduct(*mean,*mean);
  *cov /= double(nfields*(nfields-1));
  for(size_t i=0;i<ntot;i++) assert((*cov)(i,i) >= 0.);

#ifdef BMODE
  *cov /= sqrt(BMODE*BMODE+1);
#endif

  nmapout = nmap;

}

void ReadCMB(map<UParams,double>& lnlike)
{
  dbg<<"Start ReadCMB\n";

  dbg<<"Using "<<cmbfile<<" for priors\n";
  ifstream fin(cmbfile);
  if (!fin) myerror("opening cmb file:  ",cmbfile);

  double ombh2,omch2,theta,n,Z,w,As,h,ll;

  // I ignore theta and Z. 
  // ll is apparently not acually the log likelihood.
  // Licia says to ignore it and take the likelihood
  // to be proportional to the repetition of the lines.
  while (fin >> ombh2 >> omch2 >> theta >> n >> Z >> w >> As >> h >> ll) {
    UParams up(ombh2,omch2,n,w,As,h);
    lnlike[up]+=1.;  // for now, _not_ log
  }
  dbg<<"cmblnlike size = "<<lnlike.size()<<endl;
  map<UParams,double>::iterator mapit;
  for(mapit = lnlike.begin(); mapit != lnlike.end(); mapit++) {
    mapit->second = log(mapit->second);
    //dbg<<mapit->first.w<<"  "<<mapit->first.As<<"  "<<mapit->second<<endl;
  }
}

int main(int argc, char* argv[]) 
{
#ifdef NAME
  string debugname = string("predmapsq.") + NAME + ".debug";
#else
  string debugname = string("predmapsq.debug");
#endif
  dbgout = new ofstream(debugname.c_str());

  if (argc < 2) myerror("Usage: predmapsq datecode");
  string datecode = argv[1];
  dbg<<"datecode = "<<datecode<<endl;

#ifdef DILLUT
  dbg<<"DILLUT = "<<DILLUT<<endl;
#endif
#ifdef BMODE
  dbg<<"BMODE = "<<BMODE<<endl;
#endif
  dbg<<"map range = "<<map_minr<<" - "<<map_maxr<<endl;
  dbg<<"var range = "<<var_minr<<" - "<<var_maxr<<endl;


  map<UParams,double> cmblnlike;
  ReadCMB(cmblnlike);

  Vector<double>* r;
  Vector<double>* data;
  Matrix<double>* mSigma;
  size_t nmap;
  ReadData(r,data,mSigma,nmap,datecode);
  const Matrix<double>& Sigma = *mSigma;
  // This assures that the SVD doesn't get killed by Sigma(x,y) kinds of uses.
  size_t ntot = r->size();

  Sigma.DivideUsing(tmv::SV);
  Sigma.SetDiv();
  Sigma.SVD().Thresh(singthresh,dbgout);

  map<UParams,double> totlnlike;
  map<UParams,double>::iterator mapit;
#ifdef NAME
  string datname = string("lnlike.") + NAME + ".dat";
#else 
  string datname = string("lnlike.dat");
#endif
  ofstream fout(datname.c_str());
  fout <<"#  om_m    om_lam     om_b        h        As       n         w       sig8     datall      ll\n";
  int count=1;
  for(mapit=cmblnlike.begin();mapit!=cmblnlike.end();mapit++,count++) {

#ifdef STARTAT
    if (count < STARTAT) continue;
#endif

    const UParams& up = mapit->first;

    double cmbll = mapit->second;
    dbg<<count<<"/"<<cmblnlike.size()<<"  "<<up.omm<<"  "<<up.omlam<<"  "<<up.h<<"  "<<up.w<<"  "<<up.As<<endl;

    // Check for some blocks which have integration trouble and are
    // way outside the region allowed by the data:

    //if (3.*up.omlam + up.w  <= 2.1) continue;
    //if (up.omlam + up.w  <= 0.4) continue;
    //if (up.omm > 1.0) continue;

    const Universe u(up);
    //const Universe u(0.3,0.7,0.04,0.7,0.8,1.0,-1.0,0.0);
    //const Universe u(0.25,0.75,0.045,0.74,0.98,1.02,-1.0,0.0);

    double sig8 = u.sigma8();
    dbg <<"sigma8 = "<<sig8<<endl;

    // Another trouble area
    //if (log(sig8) + 0.5*log(up.omm) < -1.25) continue;
    //if (log(sig8) + 0.5*log(up.omm) > -0.3) continue;

    double datall;
    Vector<double> pred(ntot);

    for(size_t i=0;i<ntot;i++) {
      if (i < nmap) {
	pred[i] = u.fmm((*r)[i]);
      } else {
	pred[i] = u.fvar((*r)[i]);
      }
      dbg << (*r)[i]/ARCMIN <<"  "<< (*data)[i] <<"  "<< pred[i] <<"  "<< ((*data)[i]-pred[i])/sqrt(Sigma(i,i)) <<endl;
    }

    Vector<double> diff = *data - pred;
    //dbg <<"diff = "<<diff<<endl;
    //dbg <<"diff/sig = "<<(diff/Sigma)<<endl;
    //dbg <<"diff/sig*sig = "<<(diff/Sigma*Sigma)<<endl;
    double chisq = diff * (diff/Sigma);
    dbg <<"chisq = "<<chisq<<endl;
    //dbg <<"likelihood = "<<exp(-chisq/2.)<<endl;
    datall = -chisq/2.;

    double ll = datall+cmbll;
    totlnlike[up] = ll;
    dbg <<up.omm<<"  "<<up.omlam<<"  "<<up.omb<<"  "<<up.h<<"  "<<up.As<<"  "<<up.n<<"  "<<up.w<<"  "<<sig8<<"  "<<datall<<"  "<<ll<<endl;
    dbg <<"totlnlike = "<<datall<<" + "<<cmbll<<" = "<<ll<<endl;
    fout <<up.omm<<"  "<<up.omlam<<"  "<<up.omb<<"  "<<up.h<<"  "<<up.As<<"  "<<up.n<<"  "<<up.w<<"  "<<sig8<<"  "<<datall<<"  "<<ll<<endl;
  }

  return 0;
}
